package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Set;


import nlp.langmodel.LanguageModel;

public class GoodTuringBigramModel implements LanguageModel {
	private static String START = "<S>";
	private static String STOP = "</S>";
	private static double INCREMENT_FACTOR = 1.0;
	private GoodTuringCounter<Bigram<String, String>> bigramCounter = new GoodTuringCounter<Bigram<String,String>>();
	private double bigramTotal;
	
	public GoodTuringBigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		
		// bigram generation
		for (List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			
			Iterator<String> iter = stoppedSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			while (iter.hasNext()) {
				Bigram<String, String> bigram = new Bigram<String, String>(first, second);
				bigramCounter.incrementCount(bigram, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
		
		// nomarlizing counter
		bigramCounter.normalize();
		bigramTotal = bigramCounter.totalCount();
		System.out.println("bigramCount = " + bigramTotal);
		
		// train validation set
	}
	
	private double getWordLogProbability(List<String> sentence, int index) {
		if (index <= 0) {
			return 0.0;
		}
		
		String current = sentence.get(index);
		String previous = sentence.get(index - 1);
		
		Bigram<String, String> key = new Bigram<String, String>(previous, current);
		double count = bigramCounter.getCount(key) / bigramTotal;
		if (count == 0.0) {
			count = bigramCounter.getZeroCount() / bigramTotal;
		}
		return Math.log(count) / Math.log(2.0);
		
	}
	
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		double probability = 0.0;
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		
		for (int index = 0; index < stoppedSentence.size(); index++) {
			probability += getWordLogProbability(stoppedSentence, index);
		}
//		System.out.println("sentence probability is " + probability);
		return probability;
	}

	private Bigram<String, String> generateWord() {
		 double sample = Math.random();
		 double sum = 0.0;
		 for (Bigram<String, String> key : bigramCounter.keySet()) {
			 sum += bigramCounter.getCount(key) / bigramTotal;
			 if (sum >= sample) {
				 return key;
			 }
		 }
		return new Bigram<String, String>("*UNKNOWN*", STOP);
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		sentence.add(START);
		Bigram<String, String> bi = generateWord();
		while (!bi.getSecond().equals(STOP)) {
			sentence.add(bi.getFirst());
			sentence.add(bi.getSecond());
			bi = generateWord();
		}
		sentence.add(bi.getFirst());
		sentence.add(STOP);
		return sentence;
	}
}
